import platform
import re
import subprocess
from typing import List
import json

from WebBackend.dataclass.DiskItem import DiskItem, PartitionItem
from WebBackend.interface.management.Dashboard.func.disk_smart import get_disk_smart_info
from WebBackend.interface.management.Dashboard.func.disk_type import get_disk_type


def get_disk_list() -> List[DiskItem]:
    """获取系统磁盘列表（跨平台）"""
    system = platform.system()
    disks: list[DiskItem] = []

    if system == "Linux":
        try:
            # (1) 获取所有块设备信息（JSON 格式）
            cmd = ["lsblk", "-b", "-o", "NAME,TYPE,FSTYPE,SIZE,MOUNTPOINT,FSAVAIL,FSUSED,MODEL", "--json"]
            output = subprocess.check_output(cmd, universal_newlines=True, stderr=subprocess.DEVNULL)
            data = json.loads(output)

            # (2) 遍历所有设备
            for device in data.get("blockdevices", []):
                if device.get("type") != "disk":
                    continue
                # 初始化数据
                disk_info: DiskItem = {
                    "device": f"/dev/{device['name']}",
                    "model": device.get("model", "Unknown").strip(),
                    "size": round(float(device.get('size', 0)) / (1024 ** 3), 2),  # 单位GiB
                    "used": 0,
                    "type": get_disk_type(f"/dev/{device['name']}"),
                    "partitions": [],
                    'SMART': get_disk_smart_info(f"/dev/{device['name']}"),
                }
                # 情况1：检查直接挂载的分区（单分区情况）
                if device.get("mountpoint"):
                    disk_info["used"] = round(float(device.get("fsused", 0)) / (1024 ** 3), 2)  # 单位GiB
                    disk_info['partitions'] = [{
                        "name": f"/dev/{device['name']}",
                        "size": disk_info.get('size', 0),
                        "used": disk_info.get('used', 0),
                        "mount": device.get("mountpoint", 'None'),
                        "fs": device.get("fstype", "None")
                    }]
                # (3) 查找该磁盘的所有分区
                if device.get("children"):
                    for part in device["children"]:
                        if part.get("type") == "part":
                            u = part.get("fsused", 0)
                            if u is None:
                                u = 0
                            used = u / (1024 ** 3)
                            disk_info['used'] += used
                            part_info: PartitionItem = {
                                "name": f"/{part['name']}",
                                "size": round(int(part.get("size", 0)) / (1024 ** 3), 2),  # 单位GiB，保留2位小数
                                "used": used,  # 单位MiB，保留2位小数
                                "mount": part.get("mountpoint", ""),
                                "fs": part.get("fstype", "")
                            }
                            disk_info["partitions"].append(part_info)
                disks.append(disk_info)

        except (subprocess.CalledProcessError, json.JSONDecodeError):
            # 方法2：如果JSON格式失败，回退到原始方法（但改进解析）
            output = subprocess.check_output(
                ["lsblk", "-d", "-b", "-o", "NAME,MODEL,SIZE,ROTA,TYPE,VENDOR"],
                universal_newlines=True,
                stderr=subprocess.DEVNULL,
            )
            lines = output.strip().split("\n")[1:]

            for line in lines:
                # 使用正则表达式处理可能包含空格的字段
                match = re.match(
                    r"^(?P<name>\S+)\s+(?P<model>.+?)\s+(?P<size>\d+)\s+(?P<rota>\d)\s+(?P<type>\S+)\s*(?P<vendor>\S*)?$",
                    line,
                )
                if match:
                    data = match.groupdict()
                    model = data["model"] or data.get("vendor", "")
                    disk_type = "SSD" if data["rota"] == "0" else "HDD"

                    if data["type"] == "rom":
                        continue

                    disks.append(
                        {
                            "device": f"/dev/{data['name']}",
                            "model": model.strip(),
                            "size": f"{int(data['size']) // (1024 ** 3)}GB",
                            "type": (
                                "Virtual" if "virtual" in model.lower() else disk_type
                            ),
                            "interface": "Unknown",  # 原始输出中没有这个字段
                        }
                    )

        except FileNotFoundError:
            # 方法3：最终回退方案
            import glob

            for dev in glob.glob("/dev/sd?") + glob.glob("/dev/nvme?n?"):
                disks.append(
                    {
                        "device": dev,
                        "model": "Unknown",
                        "size": "Unknown",
                        "type": "Unknown",
                    }
                )

    elif system == "Darwin":
        # macOS实现
        try:
            output = subprocess.check_output(
                ["diskutil", "list"], universal_newlines=True
            )
            lines = [l for l in output.split("\n") if l.startswith("/dev/")]
            for line in lines:
                parts = line.split()
                disks.append(
                    {
                        "device": parts[0],
                        "size": parts[1],
                        "type": "SSD" if "Solid State" in line else "HDD",
                    }
                )
        except FileNotFoundError:
            pass
    elif system == "Windows":
        try:
            import wmi
            c = wmi.WMI()
            for disk in c.Win32_DiskDrive():
                disks.append(
                    {
                        "device": disk.DeviceID.replace("\\\\.\\", ""),
                        "model": disk.Model,
                        "size": (
                            f"{int(disk.Size) // (1024 ** 3)}GB"
                            if disk.Size
                            else "Unknown"
                        ),
                        "type": "SSD" if "SSD" in disk.Model.upper() else "HDD",
                        "interface": disk.InterfaceType,
                    }
                )
        except ImportError:
            raise ImportError("Please install pywin32: pip install pywin32")

    return disks


# 使用示例
if __name__ == "__main__":
    d: list[DiskItem] = get_disk_list()
    print(f"Detected {len(d)} disk(s):")
    for i, disk in enumerate(d, 1):
        print(f"\nDisk #{i}:")
        for k, v in disk.items():
            print(f"  {k:10}: {v}")
